use stencil_tysys::{
    ArgRef, BinOp, EnvExpr, FieldRef, IntEncoding, LiteralValue, LocalRef, MonOp, ValueType,
};

/// Expression with empty tags.
pub type Expr = TaggedExpr<()>;

/// Expression tagged with types.
pub type TypedExpr = TaggedExpr<ValueType>;

/// Compuable expression, possibly needing some context.  The tag can be
/// anything you want, useful for various purposes.
#[derive(Clone, Debug)]
pub enum TaggedExpr<T> {
    /// Resolve to a literal value.
    Const(T, LiteralValue),

    /// A var we read from the eval environment.
    Env(T, EnvExpr),

    /// Compute a mono-op on an expression.
    Monop(T, MonOp, Box<TaggedExpr<T>>),

    /// Compute a bin-op on an expression.
    Binop(T, BinOp, Box<TaggedExpr<T>>, Box<TaggedExpr<T>>),

    /// Evals one of two expressions dependent on an initial condition
    /// expression.
    Cond(
        T,
        Box<TaggedExpr<T>>,
        Box<TaggedExpr<T>>,
        Box<TaggedExpr<T>>,
    ),

    /// Cast a buffer to a value, if possible.
    Cast(T, ValueType, Box<TaggedExpr<T>>),

    /// Resolves to the length of a buffer.
    BufferLen(T, Box<TaggedExpr<T>>),

    /// Reads from the source buffer a value at some offset and length,
    /// performing bounds checks as necessary.  Returns a slice.
    ///
    /// (_, buffer, offset, len)
    BufferRead(
        T,
        Box<TaggedExpr<T>>,
        Box<TaggedExpr<T>>,
        Box<TaggedExpr<T>>,
    ),

    /// Reads an integer from the source buffer at some offset, performing
    /// bounds checks as necessary.  This would be used closer to how the
    /// codegen will probably be implemented under the hood than `BufferRead`.
    ///
    /// (_, buffer, offset, encoding)
    // TODO support endianness
    BufferReadInt(T, Box<TaggedExpr<T>>, Box<TaggedExpr<T>>, IntEncoding),
}

impl<T> TaggedExpr<T> {
    pub fn tag(&self) -> &T {
        match self {
            TaggedExpr::Const(t, _) => t,
            TaggedExpr::Env(t, _) => t,
            TaggedExpr::Monop(t, _, _) => t,
            TaggedExpr::Binop(t, _, _, _) => t,
            TaggedExpr::Cond(t, _, _, _) => t,
            TaggedExpr::Cast(t, _, _) => t,
            TaggedExpr::BufferLen(t, _) => t,
            TaggedExpr::BufferRead(t, _, _, _) => t,
            TaggedExpr::BufferReadInt(t, _, _, _) => t,
        }
    }

    /// Copies the expression, converting its tag according to some function.
    /// Tags for deeper nodes are mapped before shallower nodes so that it's
    /// like we're reconstructing the tree from scratch.
    pub fn map_tag<R, F: FnMut(&T) -> R>(&self, f: &mut F) -> TaggedExpr<R> {
        match self {
            Self::Const(t, a) => TaggedExpr::Const(f(t), *a),
            Self::Env(t, a) => TaggedExpr::Env(f(t), *a),
            Self::Monop(t, a, b) => {
                let bb = Box::new(b.map_tag::<R, F>(f));
                TaggedExpr::Monop(f(t), *a, bb)
            }
            Self::Binop(t, a, b, c) => {
                let bb = Box::new(b.map_tag::<R, F>(f));
                let cc = Box::new(c.map_tag::<R, F>(f));
                TaggedExpr::Binop(f(t), *a, bb, cc)
            }
            Self::Cond(t, a, b, c) => {
                let aa = Box::new(a.map_tag::<R, F>(f));
                let bb = Box::new(b.map_tag::<R, F>(f));
                let cc = Box::new(c.map_tag::<R, F>(f));
                TaggedExpr::Cond(f(t), aa, bb, cc)
            }
            Self::Cast(t, a, b) => {
                let bb = Box::new(b.map_tag::<R, F>(f));
                TaggedExpr::Cast(f(t), a.clone(), bb)
            }
            Self::BufferLen(t, a) => {
                let aa = Box::new(a.map_tag::<R, F>(f));
                TaggedExpr::BufferLen(f(t), aa)
            }
            Self::BufferRead(t, a, b, c) => {
                let aa = Box::new(a.map_tag::<R, F>(f));
                let bb = Box::new(b.map_tag::<R, F>(f));
                let cc = Box::new(c.map_tag::<R, F>(f));
                TaggedExpr::BufferRead(f(t), aa, bb, cc)
            }
            Self::BufferReadInt(t, a, b, ienc) => {
                let aa = Box::new(a.map_tag::<R, F>(f));
                let bb = Box::new(b.map_tag::<R, F>(f));
                TaggedExpr::BufferReadInt(f(t), aa, bb, *ienc)
            }
        }
    }

    /// Converts the expression to an untagged variant.
    pub fn to_untagged(&self) -> Expr {
        self.map_tag(&mut |_| ())
    }

    /// Traverses the tag tree.
    pub fn to_int_tag(&self) -> TaggedExpr<u32> {
        let mut counter = 1u32;
        let mut f = move |_: &_| {
            let v = counter;
            counter += 1;
            v
        };

        self.map_tag::<u32, _>(&mut f)
    }
}

impl<T: Clone> TaggedExpr<T> {
    /// Like `to_int_tag`, but keeps the original tag as part of a pair with the int tag.
    pub fn to_with_int_tag(&self) -> TaggedExpr<(T, u32)> {
        let mut counter = 1u32;
        let mut f = move |t: &_| {
            let tt: T = <T as Clone>::clone(t); // wtf???
            let v = counter;
            counter += 1;
            (tt, v)
        };

        self.map_tag::<(T, u32), _>(&mut f)
    }
}

/// Untagged expr constructors.
impl TaggedExpr<()> {
    /// Creates a new monop expr using an op and an expression.
    pub fn new_monop(op: MonOp, e: Expr) -> Self {
        Self::Monop((), op, Box::new(e))
    }

    /// Creates a new binop expr using an op and both sides of the expression.
    pub fn new_binop(op: BinOp, lhs: Expr, rhs: Expr) -> Self {
        Self::Binop((), op, Box::new(lhs), Box::new(rhs))
    }

    pub fn new_shl(e: Expr, n: u8) -> Self {
        Self::new_binop(BinOp::Shl, e, Self::Const((), LiteralValue::U64(n as u64)))
    }

    /// Creates a new conditional expr using a cond expr, a success expr, and a
    /// failure expr.
    pub fn new_cond(c: Expr, a: Expr, b: Expr) -> Self {
        Self::Cond((), Box::new(c), Box::new(a), Box::new(b))
    }

    pub fn new_read_int(buf: Expr, off: Expr, enc: IntEncoding) -> Self {
        Self::BufferReadInt((), Box::new(buf), Box::new(off), enc)
    }

    pub fn new_read_u8_at(buf: Expr, off: u64) -> Self {
        Self::new_read_int(buf, Self::new_usize(off), IntEncoding::byte())
    }

    pub fn new_cast(ty: ValueType, e: Expr) -> Self {
        Self::Cast((), ty, Box::new(e))
    }

    /// Converts some value into its literal form.
    pub fn new_const<T: Into<LiteralValue>>(value: T) -> Self {
        Self::Const((), value.into())
    }

    pub fn new_usize(v: u64) -> Self {
        Self::Const((), LiteralValue::USize(v))
    }

    pub fn new_field(field: FieldRef) -> Self {
        Self::Env((), EnvExpr::Field(field))
    }

    pub fn new_arg(arg: ArgRef) -> Self {
        Self::Env((), EnvExpr::Arg(arg))
    }

    pub fn new_local(local: LocalRef) -> Self {
        Self::Env((), EnvExpr::Local(local))
    }

    pub fn new_void() -> Self {
        Self::new_const(())
    }

    pub fn new_eq(a: Expr, b: Expr) -> Self {
        Self::new_binop(BinOp::Eq, a, b)
    }

    pub fn new_ne(a: Expr, b: Expr) -> Self {
        Self::new_binop(BinOp::Ne, a, b)
    }
}

impl From<LiteralValue> for Expr {
    fn from(value: LiteralValue) -> Self {
        Self::Const((), value)
    }
}

/// Typed expr constructors.
impl TypedExpr {
    pub fn new_read_int(buf: TypedExpr, off: TypedExpr, enc: IntEncoding) -> Self {
        Self::BufferReadInt(enc.ty().into(), Box::new(buf), Box::new(off), enc)
    }

    /// Converts some value into its literal form.
    pub fn new_const<T: Into<LiteralValue>>(value: T) -> Self {
        let lv = value.into();
        Self::Const(lv.ty(), lv)
    }

    pub fn new_void() -> Self {
        Self::new_const(())
    }

    pub fn new_usize(v: u64) -> Self {
        Self::new_const(LiteralValue::USize(v))
    }

    pub fn new_cast(ty: ValueType, v: TypedExpr) -> Self {
        Self::Cast(ty.clone(), ty, Box::new(v))
    }
}

#[cfg(test)]
mod tests {
    use crate::Expr;
    use stencil_tysys::ValueType;

    use super::*;

    fn get_012_cond_expr(bufe: Expr) -> Expr {
        let e1 = Expr::new_read_u8_at(bufe.clone(), 0);
        let e1cast = Expr::new_cast(ValueType::USize, e1);
        let e2 = Expr::new_read_u8_at(bufe.clone(), 1);
        let e3 = Expr::new_read_u8_at(bufe, 2);
        let ec = Expr::new_binop(BinOp::Eq, e1cast, Expr::new_usize(2));
        let e = Expr::new_cond(ec, e2, e3);
        e
    }

    #[test]
    fn test_int_tag() {
        let e = get_012_cond_expr(Expr::new_field(FieldRef::from(0 as u16)));
        let et = e.to_int_tag().map_tag(&mut |t| 100 * t);
        eprintln!("{et:#?}");
        assert_eq!(*et.tag(), 1300, "tag not traversed correctly")
    }
}
